#include "sample.h"
#include <iostream>
#define EIGEN_USE_THREADS
#define EIGEN_TEST_NO_LONGDOUBLE
#define EIGEN_TEST_NO_COMPLEX
#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
#include <stdio.h>
#include <unsupported/Eigen/CXX11/Tensor>
#define N 8

using namespace Eigen;

typedef Eigen::DenseIndex IndexType;
typedef Eigen::ThreadPoolDevice CPUDevice;
#define NDIMS 3
template<typename T>
using TTensor =  Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using Flat = Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using  ConstScalar = Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>,Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using ConstFlat =  Eigen::TensorMap<
Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned>;
const int esp = 1e5 + 1;

template<typename T>
struct test_cpu{
  void call(T* _m,T* _v,T* _var,T* _grad,T* f1, T*f2,T*f3,T*f4,T*f5) {
    ThreadPool p(4);
    CPUDevice d(&p,4);
    Flat<T> m(_m,N),v(_v,N),var(_var,N);
    ConstScalar<T> beta1_power(f1), lr(f2),beta1(f3),
      beta2(f4),epsilon(f5);
    ConstFlat<T> grad(_grad,N);
    m.device(d) += (grad - m)*(T(1) - beta1());
    v.device(d) = (beta2()*v).cwiseMax(grad.abs());
    var.device(d) -= lr() / (T(1) - beta1_power()) * (m/(v + epsilon()));
    auto tmp = lr();
    printf("cpu: %f\n", _var[0]);
  }
};
int eigen_cpu_sample() {
  float _v[] = {0.3042 ,0.525  ,0.432  ,0.2913 ,0.612  ,0.1395 ,0.2922 ,0.3665};
  float _m[] = {0.601   ,0.708   ,0.02058 ,0.9697  ,0.8325  ,0.2123  ,0.1819  ,0.1833};
  float _grad[] = {0.456   ,0.785   ,0.1997  ,0.514   ,0.5923  ,0.04645 ,0.6074  ,0.1705 };
  float _var[] = {0.3745  ,0.9507  ,0.732   ,0.5986  ,0.156   ,0.156   ,0.05807 ,0.866 };
  float beta1_power = 0.9;
  float lr = 0.01;
  float beta1 = 0.9;
  float beta2 = 0.99;
  float epsilon = 1e-8;
  float *_m_dev, *_v_dev, *_var_dev, *_grad_dev, *beta1_power_dev,
        *lr_dev, *beta1_dev, *beta2_dev, *epsilon_dev;

  test_cpu<float> cpu_gen;
  cpu_gen.call(_m, _v, _var, _grad,
      &beta1_power, &lr, &beta1, &beta2, &epsilon);
  return 0;
}
